import json
from transformers import AutoTokenizer
from tqdm import tqdm  # 进度条工具


def analyze_qwen_prompt_tokens(json_file_path: str):
    """
    分析JSON文件中prompt字段的token数量（针对Qwen3-8B模型）
    :param json_file_path: JSON文件路径
    """
    # 加载Qwen3-8B的tokenizer
    print("正在加载Qwen3-8B tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B", trust_remote_code=True)

    # 初始化统计变量
    stats = {
        "total_items": 0,
        "max_tokens": 0,
        "min_tokens": float("inf"),
        "avg_tokens": 0.0,
        "over_2048": 0,
        "over_3072": 0,
        "over_4096": 0
    }
    total_tokens = 0

    try:
        # 读取JSON文件
        print(f"正在读取JSON文件: {json_file_path}")
        with open(json_file_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        print(f"找到 {len(data)} 条记录，开始分析token数量...")

        # 使用tqdm创建进度条
        for item in tqdm(data, desc="分析进度"):
            prompt = item.get("prompt", "")
            if not isinstance(prompt, str):
                continue

            # 使用Qwen tokenizer计算token数量
            tokens = tokenizer.encode(prompt, add_special_tokens=False)
            token_count = len(tokens)

            # 更新统计信息
            stats["total_items"] += 1
            total_tokens += token_count

            if token_count > stats["max_tokens"]:
                stats["max_tokens"] = token_count
            if token_count < stats["min_tokens"]:
                stats["min_tokens"] = token_count

            if token_count > 2048:
                stats["over_2048"] += 1
            if token_count > 3072:
                stats["over_3072"] += 1
            if token_count > 4096:
                stats["over_4096"] += 1

        # 计算平均值
        if stats["total_items"] > 0:
            stats["avg_tokens"] = round(total_tokens / stats["total_items"], 2)
        else:
            stats["min_tokens"] = 0  # 没有数据时重置最小值

        return stats

    except Exception as e:
        print(f"处理过程中发生错误: {e}")
        return None


def print_stats(stats: dict):
    """打印统计结果"""
    if not stats:
        return

    print("\n===== Qwen3-8B Token统计结果 =====")
    print(f"总记录数: {stats['total_items']}")
    print(f"Token最大值: {stats['max_tokens']}")
    print(f"Token最小值: {stats['min_tokens']}")
    print(f"Token平均值: {stats['avg_tokens']}")

    if stats['total_items'] > 0:
        print(
            f"Token超过2048的记录数: {stats['over_2048']} (占比: {stats['over_2048'] / stats['total_items'] * 100:.2f}%)")
        print(
            f"Token超过3072的记录数: {stats['over_3072']} (占比: {stats['over_3072'] / stats['total_items'] * 100:.2f}%)")
        print(
            f"Token超过4096的记录数: {stats['over_4096']} (占比: {stats['over_4096'] / stats['total_items'] * 100:.2f}%)")
    else:
        print("没有可用的记录")


if __name__ == "__main__":
    # 替换为您的JSON文件路径
    json_file = "filtered_data1.json"

    # 分析token数量
    stats = analyze_qwen_prompt_tokens(json_file)

    if stats:
        print_stats(stats)
